# -*- coding: utf-8 -*-
"""
@Time    : 2025/2/17 10:58 
@Author  : ZhangShenao 
@File    : 7.tool_calling.py 
@Desc    : 工具调用
"""
import json
import os

import dotenv
from zhipuai import ZhipuAI

from convert_function_to_tool import convert_function_to_tool


def get_current_weather(location: str, unit: str = "fahrenheit") -> float:
    """获取当前天气信息"""
    if unit == "celsius":
        return 25.0
    else:
        return 77.0


if __name__ == '__main__':
    # 加载环境变量
    dotenv.load_dotenv()

    # 创建ZhipuAI客户端
    client = ZhipuAI(api_key=os.getenv("ZHIPUAI_API_KEY"))

    # 将函数转换为工具调用参数
    weather_tool = convert_function_to_tool(get_current_weather)
    tools = [weather_tool]

    # 在内存中保存上下文
    messages = []

    # 循环接收用户输入
    while True:
        query = input("User: ")
        if query.lower() == "bye":
            print("bye bye~")
            break

        message = {"role": "user", "content": query}
        messages.append(message)
        response = client.chat.completions.create(
            model="glm-4-flash",
            messages=messages,
            tools=tools,
            tool_choice="auto"
        ).choices[0]

        # 普通消息
        if response.message.tool_calls is None:
            print("AI: ", response.message.content)
            messages.append({"role": "assistant", "content": response.message.content})
        # 工具消息
        else:
            tool_message = response.message.tool_calls[0]
            args = tool_message.function.arguments

            # 调用工具,获取结果
            result = get_current_weather(**json.loads(args))
            messages.append({"role": "tool", "content": f"{result}"})

            # 将工具调用结果发送给大模型,生成最终结果
            response = client.chat.completions.create(
                model="glm-4-flash",
                messages=messages,
                tools=tools,
                tool_choice="auto"
            ).choices[0]
            print("AI: ", response.message.content)
